import socket
import struct

from PIL import Image

class DataParser():
    data = None
    offset = 0
    data_len = 0
	
    def parse_small_int(self):
        if not self.check_bounds(1):
            return None

        val = ord(self.data[self.offset])
        self.offset += 1
        return val

    def parse_int(self):
        if not self.check_bounds(4):
            return None
        val = struct.unpack_from(">i", self.data, self.offset)[0]
        self.offset += 4
        return val

    def parse_string(self):
        if not self.check_bounds(2):
            return None
        length = struct.unpack_from(">i", "\0\0"+self.data[self.offset:self.offset+2])[0]
        self.offset += 2

        if not self.check_bounds(length):
            return None
        data = self.data[self.offset:self.offset+length]
        self.offset += length

        return str(data)

    def check_bounds(self, num):
        return self.data_len >= self.offset + num

    def parse_list(self):
        if not self.check_bounds(5):
            return None

        length = struct.unpack_from(">i", self.data, self.offset)[0]
        self.offset = 6

        elements = []

        for i in range(0, length):
            type = ord(self.data[self.offset])
            self.offset += 1

            if type == 97:
                element = self.parse_small_int()
            elif type == 98:
                element = self.parse_int()
            elif type == 107:
                element = self.parse_string()
            else:
                raise Exception("Unknown data type "+str(type)+" in erlang bynary list")

            if not element is None:
                elements.append(element)
            else:
                return None

        if not self.check_bounds(1):
            return None

        self.offset += 1 # Skip the NIL element.
        return tuple(elements)

    def parse_byte_list(self):
        if not self.check_bounds(2):
            return None

        length = struct.unpack_from(">i", "\0\0"+self.data[self.offset:self.offset+2])[0]
        self.offset += 2

        if not self.check_bounds(length):
            return None

        elements = []
        for i in range(0, length):
            elements.append(self.parse_small_int())

        return tuple(elements)

    def get_remaining_data(self):
        return self.data[self.offset:]

    def get_term(self, data):
        self.data = data
        self.data_len = len(self.data)
        self.offset = 0

        ver = ord(self.data[0])
        if ver == 131:	# known version
            type = ord(self.data[1])
            self.offset = 2

            if type == 108:	# list
                return self.parse_list()
            elif type == 107:	# byte string
                return self.parse_byte_list()
            else:
                raise Exception("Unknown data type "+str(type)+" in erlang binary")
        else:
            self.offset = 0
            return None
		
def run(image, socket):
    data_parser = DataParser()

    socket.listen(1)
    connection, address = socket.accept()

    data = r""
    while not image.is_ready():
        buff = connection.recv(65536)
        data += buff

        while data:
            term = data_parser.get_term(data)

            if term:
                cmd(term, image)
                data = data_parser.get_remaining_data()
            else:
                break

    image.save()

		
def cmd(c, image):
    print c
    if c[0] == "px":
        c, x, y, r, g, b = c
        image.set_pixel(x, y, (r, g, b))

    elif c[0] == "set_size":
        w = c[1]
        h = c[2]
        image.create(w, h)

    elif c[0] == "exit":
        image.save()
	
class RenderImage():
    _img = None
    _pix = None
    _w = 0
    _h = 0
    _buff = []
    _num = 0
    _save_after = 0

    def is_ready(self):
        if self._img is None:
            return False

        if self._num > 0:
            return False
        else:
            return True

    def create(self, width, height):
        self._w = width
        self._h = height
        self._img = Image.new("RGB", (width, height))
        self._pix = self._img.load()

        self._num = width*height*10000000 #The image should never be ready, so this should be really big (temporary, hack).

        self._save_after = width*10

        if self._buff:
            for pixel in self._buff:
                self.set_pixel(pixel[0], pixel[1], pixel[2])

        print "Image created"

    def set_pixel(self, x, y, color):
        if self._pix is None:
            self._buff.append((x, y, color))
            return

        print str(self._num)
        if self._num % self._save_after == 0:
            self.save()

        self._num -= 1
        r, g, b = color
        try:
            self._pix[x, y] = (r, g, b)
        except IndexError:
            self.save()
            print "Index error at {x},{y}".format(x=x, y=y)

    def save(self):
        self._img.save("render.png", "PNG")
        print "Image saved"


def main():
    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    s.bind(('localhost', 5621))

    image = RenderImage()
    try:
        run(image, s)
    except:
        raise
    finally:
        s.close()

	
if __name__ == '__main__':
    main()
